import tensorflow as tf
import numpy as np
import pdb


def conv(x, filter_height, filter_width, num_filters,
         stride_x, stride_y, name, padding='SAME', groups=1):
    # 数据的最后一个维度是channel的数量
    input_channels = int(x.get_shape()[-1])
    def convolve(i, k): return tf.nn.conv2d(i, k, strides=[
        1, stride_y, stride_x, 1], padding=padding)
    with tf.variable_scope(name) as scope:
        weights = tf.get_variable('weights', shape=[filter_height,
                                                    filter_width,
                                                    input_channels/groups,
                                                    num_filters])
        biases = tf.get_variable('biases', shape=[num_filters])
    if groups == 1:
        conv = convolve(x, weights)
    add = tf.nn.bias_add(conv, biases)
    bias = tf.reshape(add, tf.shape(conv))
    relu = tf.nn.relu(bias, name=scope.name)
    return relu


def fc(x, num_in, num_out, name, relu=True):
    with tf.variable_scope(name) as scope:
        weights = tf.get_variable('weights', shape=[num_in, num_out],
                                  trainable=True)
        biases = tf.get_variable('biases', [num_out], trainable=True)
        act = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)
    if relu:
        relu = tf.nn.relu(act)
        return relu
    else:
        return act


def max_pool(x, filter_height, filter_width, stride_y, stride_x, name,
             padding='SAME'):
    return tf.nn.max_pool(x, ksize=[1, filter_height, filter_width, 1],
                          strides=[1, stride_y, stride_x, 1],
                          padding=padding, name=name)


def lrn(x, radius, alpha, beta, name, bias=1.0):
    return tf.nn.local_response_normalization(x, depth_radius=radius,
                                              alpha=alpha, beta=beta, bias=bias, name=name)


def dropout(x, keep_prob):
    return tf.nn.dropout(x, keep_prob)


class AlexNet(object):

    def __init__(self, x, keep_prob, num_classes, skip_layer,
                 weights_path='DEFAULT'):
        self.X = x
        self.NUM_CLASSES = num_classes
        self.KEEP_PROB = keep_prob
        self.SKIP_LAYER = skip_layer
        if weights_path == 'DEFAULT':
            self.WEIGHTS_PATH = '../bvlc_alexnet.npy'
        else:
            self.WEIGHTS_PATH = weights_path

        self.create()

    def create(self):
        conv1 = conv(self.X, 11, 11, 96, 4, 4, padding='VALID', name='conv1')
        norm1 = lrn(conv1, 2, 2e-05, 0.75, name='norm1')
        pool1 = max_pool(norm1, 3, 3, 2, 2, padding='VALID', name='pool1')

        conv2 = conv(pool1, 5, 5, 256, 1, 1, name='conv2')
        norm2 = lrn(conv2, 2, 2e-04, 0.75, name='norm2')
        pool2 = max_pool(norm2, 3, 3, 2, 2, padding='VALID', name='pool2')

        conv3 = conv(pool2, 3, 3, 384, 1, 1, name='conv3')
        conv4 = conv(conv3, 3, 3, 384, 1, 1, name='conv4')
        conv5 = conv(conv4, 3, 3, 256, 1, 1, name='conv5')
        pool5 = max_pool(conv5, 3, 3, 2, 2, padding='VALID', name='pool5')
        # pdb.set_trace()
        flattend = tf.reshape(pool5, [-1, 6*6*256])
        fc6 = fc(flattend, 6*6*256, 4096, name='fc6')
        dropout6 = dropout(fc6, self.KEEP_PROB)

        fc7 = fc(dropout6, 4096, 4096, name='fc7')

        dropout7 = dropout(fc7, self.KEEP_PROB)

        self.fc8 = fc(dropout7, 4096, self.NUM_CLASSES, relu=False, name='fc8')

    def load_initial_weights(self, session):
        weights_dict = np.load(self.WEIGHTS_PATH, encoding='bytes').item()
        for name in weights_dict:
            print(name)
        for op_name in weights_dict:
            if op_name not in self.SKIP_LAYER:
                with tf.variable_scope(op_name, reuse=True):
                    for data in weights_dict[op_name]:
                        if len(data.shape) == 1:
                            var = tf.get_variable(
                                'biases', trainable=False)
                            print(var)
                            print(type(data))
                            pdb.set_trace()

                            session.run(var.assign(data))
                        else:
                            var = tf.get_variable(
                                'weights', trainable=False)
                            session.run(var.assign(data))
                            print(var)
                            print(type(data))
                            pdb.set_trace()
